// WebSocket.cpp
#include "WebSocket.hpp"
#include <openssl/sha.h>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/buffer.h>
#include <cstring>

bool WebSocket::handleHandshake(const std::string& request, std::string& response) {
    // 查找 Sec-WebSocket-Key
    size_t keyPos = request.find("Sec-WebSocket-Key: ");
    if (keyPos == std::string::npos) {
        return false;
    }
    
    size_t keyStart = keyPos + 19; // "Sec-WebSocket-Key: " 的长度
    size_t keyEnd = request.find("\r\n", keyStart);
    if (keyEnd == std::string::npos) {
        return false;
    }
    
    std::string key = request.substr(keyStart, keyEnd - keyStart);
    
    // 生成 Sec-WebSocket-Accept
    std::string acceptKey = generateAcceptKey(key);
    
    // 构造响应
    response = "HTTP/1.1 101 Switching Protocols\r\n"
               "Upgrade: websocket\r\n"
               "Connection: Upgrade\r\n"
               "Sec-WebSocket-Accept: " + acceptKey + "\r\n\r\n";
    
    return true;
}

std::string WebSocket::encodeFrame(const std::string& message, Opcode opcode) {
    std::string frame;
    
    // 第一个字节：FIN + RSV1-3 + opcode
    uint8_t firstByte = 0x80 | (opcode & 0x0F);
    frame.push_back(firstByte);
    
    // 第二个字节：MASK + Payload length
    size_t payloadLength = message.length();
    uint8_t secondByte = 0x00; // 不使用掩码（服务器到客户端）
    
    if (payloadLength <= 125) {
        secondByte |= payloadLength;
        frame.push_back(secondByte);
    } else if (payloadLength <= 65535) {
        secondByte |= 126;
        frame.push_back(secondByte);
        frame.push_back((payloadLength >> 8) & 0xFF);
        frame.push_back(payloadLength & 0xFF);
    } else {
        secondByte |= 127;
        frame.push_back(secondByte);
        for (int i = 7; i >= 0; i--) {
            frame.push_back((payloadLength >> (8 * i)) & 0xFF);
        }
    }
    
    // Payload data
    frame.append(message);
    
    return frame;
}

bool WebSocket::decodeFrame(const char* data, size_t len, std::string& message, Opcode& opcode) {
    if (len < 2) return false;
    
    const uint8_t* bytes = reinterpret_cast<const uint8_t*>(data);
    
    // 解析第一个字节
    bool fin = bytes[0] & 0x80;
    opcode = static_cast<Opcode>(bytes[0] & 0x0F);
    
    // 解析第二个字节
    bool mask = bytes[1] & 0x80;
    uint64_t payloadLength = bytes[1] & 0x7F;
    
    size_t offset = 2;
    
    // 处理扩展长度字段
    if (payloadLength == 126) {
        if (len < 4) return false;
        payloadLength = (bytes[2] << 8) | bytes[3];
        offset = 4;
    } else if (payloadLength == 127) {
        if (len < 10) return false;
        payloadLength = 0;
        for (int i = 0; i < 8; i++) {
            payloadLength = (payloadLength << 8) | bytes[2 + i];
        }
        offset = 10;
    }
    
    // 处理掩码密钥
    uint8_t maskingKey[4];
    if (mask) {
        if (len < offset + 4) return false;
        memcpy(maskingKey, bytes + offset, 4);
        offset += 4;
    }
    
    // 检查是否有足够的数据
    if (len < offset + payloadLength) return false;
    
    // 解码消息
    message.resize(payloadLength);
    if (mask) {
        for (size_t i = 0; i < payloadLength; i++) {
            message[i] = data[offset + i] ^ maskingKey[i % 4];
        }
    } else {
        message.assign(data + offset, payloadLength);
    }
    
    return true;
}

std::string WebSocket::generateAcceptKey(const std::string& key) {
    std::string magicString = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    std::string concatenated = key + magicString;
    
    unsigned char hash[SHA_DIGEST_LENGTH];
    SHA1(reinterpret_cast<const unsigned char*>(concatenated.c_str()), concatenated.length(), hash);
    
    // Base64 编码
    return base64Encode(std::string(reinterpret_cast<char*>(hash), SHA_DIGEST_LENGTH));
}

std::string WebSocket::base64Encode(const std::string& data) {
    BIO *bio, *b64;
    BUF_MEM *bufferPtr;

    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new(BIO_s_mem());
    bio = BIO_push(b64, bio);

    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
    BIO_write(bio, data.c_str(), data.length());
    BIO_flush(bio);
    BIO_get_mem_ptr(bio, &bufferPtr);
    
    std::string result(bufferPtr->data, bufferPtr->length);
    
    BIO_free_all(bio);
    return result;
}